class FeatureExtractor:
    def __init__(self, target_fs=12000):
        self.target_fs = target_fs
        self.feature_names = []
        
    def extract_all_features(self, signal, rpm, bearing_type='SKF6205'):
        """提取所有类型的特征"""
        features = {}
        
        # 时域特征
        time_features = self._extract_time_domain_features(signal)
        features.update(time_features)
        
        # 频域特征
        freq_features = self._extract_frequency_domain_features(signal, self.target_fs)
        features.update(freq_features)
        
        # 时频域特征
        time_freq_features = self._extract_time_frequency_features(signal, self.target_fs)
        features.update(time_freq_features)
        
        # 轴承故障特征频率
        fault_freq_features = self._extract_fault_frequency_features(signal, self.target_fs, rpm, bearing_type)
        features.update(fault_freq_features)
        
        # 非线性特征
        nonlinear_features = self._extract_nonlinear_features(signal)
        features.update(nonlinear_features)
        
        return features
    
    def _extract_time_domain_features(self, signal):
        """提取时域特征"""
        features = {}
        
        # 基本统计特征
        features['mean'] = np.mean(signal)
        features['std'] = np.std(signal)
        features['max'] = np.max(signal)
        features['min'] = np.min(signal)
        features['peak_to_peak'] = np.ptp(signal)
        features['rms'] = np.sqrt(np.mean(signal**2))
        
        # 形状特征
        features['skewness'] = self._calculate_skewness(signal)
        features['kurtosis'] = self._calculate_kurtosis(signal)
        features['crest_factor'] = np.max(np.abs(signal)) / features['rms']
        features['shape_factor'] = features['rms'] / np.mean(np.abs(signal))
        features['impulse_factor'] = np.max(np.abs(signal)) / np.mean(np.abs(signal))
        
        # 其他时域特征
        features['clearance_factor'] = np.max(np.abs(signal)) / (np.mean(np.sqrt(np.abs(signal)))**2)
        
        return features
    
    def _extract_frequency_domain_features(self, signal, fs):
        """提取频域特征"""
        features = {}
        
        # FFT变换
        fft_signal = np.fft.fft(signal)
        freq = np.fft.fftfreq(len(signal), 1/fs)
        positive_freq = freq[:len(freq)//2]
        positive_fft = np.abs(fft_signal[:len(fft_signal)//2])
        
        # 频域统计特征
        features['freq_mean'] = np.mean(positive_fft)
        features['freq_std'] = np.std(positive_fft)
        features['freq_max'] = np.max(positive_fft)
        features['freq_centroid'] = np.sum(positive_freq * positive_fft) / np.sum(positive_fft)
        
        # 频带能量特征
        bands = [(0, fs/8), (fs/8, fs/4), (fs/4, fs/2)]
        for i, (low, high) in enumerate(bands):
            mask = (positive_freq >= low) & (positive_freq <= high)
            band_energy = np.sum(positive_fft[mask]**2)
            features[f'band_{i+1}_energy'] = band_energy
        
        return features
    
    def _extract_time_frequency_features(self, signal, fs):
        """提取时频域特征"""
        features = {}
        
        # 小波变换特征
        coeffs = pywt.wavedec(signal, 'db4', level=4)
        for i, coeff in enumerate(coeffs):
            features[f'wavelet_{i}_energy'] = np.sum(coeff**2)
            features[f'wavelet_{i}_std'] = np.std(coeff)
        
        # STFT特征
        f, t, Zxx = signal.stft(signal, fs, nperseg=256)
        stft_magnitude = np.abs(Zxx)
        features['stft_mean'] = np.mean(stft_magnitude)
        features['stft_std'] = np.std(stft_magnitude)
        
        return features
    
    def _extract_fault_frequency_features(self, signal, fs, rpm, bearing_type):
        """提取故障特征频率相关的特征"""
        features = {}
        
        # 计算轴承参数
        bearing_params = self._get_bearing_parameters(bearing_type)
        fr = rpm / 60  # 转频
        
        # 计算故障特征频率
        bpfo = fr * bearing_params['n_balls'] / 2 * (1 - bearing_params['d']/bearing_params['D'])
        bpfi = fr * bearing_params['n_balls'] / 2 * (1 + bearing_params['d']/bearing_params['D'])
        bsf = fr * bearing_params['D']/bearing_params['d'] * (1 - (bearing_params['d']/bearing_params['D'])**2)
        ftf = fr / 2 * (1 - bearing_params['d']/bearing_params['D'])
        
        # 计算故障频率附近的能量
        fft_signal = np.fft.fft(signal)
        freq = np.fft.fftfreq(len(signal), 1/fs)
        
        for freq_type, center_freq in [('bpfo', bpfo), ('bpfi', bpfi), ('bsf', bsf), ('ftf', ftf)]:
            if center_freq > 0:
                # 计算谐波
                for harmonic in range(1, 4):
                    target_freq = center_freq * harmonic
                    bandwidth = target_freq * 0.1  # 10%带宽
                    mask = (np.abs(freq) >= target_freq - bandwidth) & (np.abs(freq) <= target_freq + bandwidth)
                    energy = np.sum(np.abs(fft_signal[mask])**2)
                    features[f'{freq_type}_harmonic_{harmonic}_energy'] = energy
        
        return features
    
    def _extract_nonlinear_features(self, signal):
        """提取非线性特征"""
        features = {}
        
        # 近似熵
        features['approx_entropy'] = self._calculate_approximate_entropy(signal)
        
        # 样本熵
        features['sample_entropy'] = self._calculate_sample_entropy(signal)
        
        # 排列熵
        features['permutation_entropy'] = self._calculate_permutation_entropy(signal)
        
        return features
    
    def _get_bearing_parameters(self, bearing_type):
        """获取轴承参数"""
        params = {
            'SKF6205': {'d': 0.3126, 'D': 1.537, 'n_balls': 9},
            'SKF6203': {'d': 0.2656, 'D': 1.122, 'n_balls': 9}
        }
        return params.get(bearing_type, params['SKF6205'])
    
    def _calculate_skewness(self, signal):
        """计算偏度"""
        return np.mean((signal - np.mean(signal))**3) / (np.std(signal)**3)
    
    def _calculate_kurtosis(self, signal):
        """计算峭度"""
        return np.mean((signal - np.mean(signal))**4) / (np.std(signal)**4)
    
    def _calculate_approximate_entropy(self, signal, m=2, r=0.2):
        """计算近似熵"""
        # 实现略
        return 0.5
    
    def _calculate_sample_entropy(self, signal, m=2, r=0.2):
        """计算样本熵"""
        # 实现略
        return 0.6
    
    def _calculate_permutation_entropy(self, signal, m=3, delay=1):
        """计算排列熵"""
        # 实现略
        return 0.7

# 批量提取特征
feature_extractor = FeatureExtractor()
all_features = []

for data in tqdm(all_data, desc="Extracting features"):
    features = feature_extractor.extract_all_features(
        data['de_signal'], data['rpm'], 'SKF6205'
    )
    features['file_name'] = data['file_name']
    features['fault_type'] = data['fault_type']
    features['fault_size'] = data['fault_size']
    features['load_condition'] = data['load_condition']
    all_features.append(features)

features_df = pd.DataFrame(all_features)
